# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import unittest
import logging

from coin.base.timestamp import get_timestamp

from coin.exchange.util.rate_limit import WindowLimiter
from coin.exchange.bitflyer_v1.order_gateway import BitflyerOrderGateway


class WindowLimiterListTest(object):
  def __init__(self, window_size_ns, transact_limit):
    self._window_size_ns = window_size_ns
    self._transact_limit = transact_limit
    self._trans_times = []

  def get_transact_size(self, timestamp):
    # assume ascending transaction timestamp
    for i, trans_time in enumerate(self._trans_times):
      if trans_time >= timestamp - self._window_size_ns:
        self._trans_times = self._trans_times[i:]
        break
    return len(self._trans_times)

  def _get_tps_approval(self, timestamp):
    tsize = self.get_transact_size(timestamp)
    if tsize < self._transact_limit:
      self._trans_times.append(timestamp)
      return True
    else:
      return False


class WindowLimiterDequeTest(object):
  def __init__(self, window_size_ns, transact_limit):
    self._window_size_ns = window_size_ns
    self._transact_limit = transact_limit
    self._trans_times = collections.deque()

  def get_transact_size(self, timestamp):
    # assume ascending transaction timestamp
    max_time = timestamp - self._window_size_ns
    while self._trans_times and self._trans_times[0] <= max_time:
      self._trans_times.popleft()
    return len(self._trans_times)

  def _get_tps_approval(self, timestamp):
    tsize = self.get_transact_size(timestamp)
    if tsize < self._transact_limit:
      self._trans_times.append(timestamp)
      return True
    else:
      return False


class EmptyOg(object):
  def cancel_multiple(*args):
    pass


class BitflyerOGForTest(BitflyerOrderGateway):
  def __init__(self, tps_limit):
    self._tps_limit = tps_limit
    self._logger = logging.getLogger("Test")


class TpsLimitTest(unittest.TestCase):
  def test_tps_safe(self):
    tps_limit_safe = WindowLimiter(5, 2)
    self.assertTrue(tps_limit_safe._get_tps_approval(1))
    self.assertTrue(tps_limit_safe._get_tps_approval(2))
    self.assertFalse(tps_limit_safe._get_tps_approval(3))
    self.assertFalse(tps_limit_safe._get_tps_approval(5))
    self.assertTrue(tps_limit_safe._get_tps_approval(7))

    # use cached.so 1 msg latent
    traits = BitflyerOGForTest(WindowLimiter(5, 2))
    self.assertTrue(traits.get_tps_approval(1))
    self.assertTrue(traits.get_tps_approval(2))
    self.assertFalse(traits.get_tps_approval(3))
    self.assertFalse(traits.get_tps_approval(5))
    self.assertTrue(traits.get_tps_approval(7))

  def test_tps_speed(self):
    tps_limit_live = WindowLimiter(5000, 500)
    tps_limit_deque = WindowLimiterDequeTest(5000, 500)
    tps_limit_pylist = WindowLimiterListTest(5000, 500)
    for tps_limit in [tps_limit_live, tps_limit_deque, tps_limit_pylist]:
      begin = get_timestamp()
      print("begin")
      ncall = 10000
      for i in range(ncall):
        tps_limit._get_tps_approval(i)
      elapsedtime = (get_timestamp() - begin) * 1e-9
      print("elap %s percall %s ns" % (elapsedtime, elapsedtime / ncall * 1e9))


if __name__ == "__main__":
  unittest.main()
